import asyncio  
import json  
import time  
from pynput.keyboard import Controller, Key  
from quart import Quart, websocket  
  
app = Quart(__name__)  
keyboard = Controller()  
  
# 支持的按键映射 - 专注于 x 和 z  
SUPPORTED_KEYS = {  
    'x': 'x',  
    'z': 'z'  
}  
  
# 存储按键状态  
key_states = {  
    'x': False,  
    'z': False  
}  
  
# 存储活跃的 WebSocket 连接  
active_connections = set()  
  
# WebSocket 连接处理  
@app.websocket("/ws")  
async def ws():  
    # 将当前的websocket对象添加到活跃连接中  
    active_connections.add(websocket._get_current_object())  
      
    try:  
        while True:  
            data = await websocket.receive()  
            await process_message(data)  
    except asyncio.CancelledError:  
        pass  
    finally:  
        active_connections.remove(websocket._get_current_object())  
  
# 处理接收到的 WebSocket 消息  
async def process_message(message):  
    try:  
        data = json.loads(message)  
        action = data.get('action')  
        key = data.get('key', '').lower()  
          
        if key not in SUPPORTED_KEYS and action != 'reset_all' and action != 'get_state':  
            await websocket.send(json.dumps({  
                "status": "error",  
                "message": f"Unsupported key: {key}"  
            }))  
            return  
              
        start_time = time.time()  
          
        if action == 'press_down':  
            # 先释放所有其他按下的键  
            for other_key in key_states:  
                if other_key != key and key_states[other_key]:  
                    key_states[other_key] = False  
                    await asyncio.to_thread(keyboard.release, SUPPORTED_KEYS[other_key])  
                    print(f"[Direct Input] Auto-released key: {other_key}")  
              
            # 然后按下新键  
            key_states[key] = True  
            await asyncio.to_thread(keyboard.press, SUPPORTED_KEYS[key])  
            print(f"[Direct Input] Key Down: {key}")  
              
            latency = (time.time() - start_time) * 1000  
            await websocket.send(json.dumps({  
                "status": "success",  
                "action": "key_down",  
                "key": key,  
                "latency": latency,  
                "key_states": key_states  
            }))  
                  
        elif action == 'press_up':  
            # 总是释放按键，无论当前状态如何  
            key_states[key] = False  
            await asyncio.to_thread(keyboard.release, SUPPORTED_KEYS[key])  
            print(f"[Direct Input] Key Up: {key}")  
              
            latency = (time.time() - start_time) * 1000  
            await websocket.send(json.dumps({  
                "status": "success",  
                "action": "key_up",  
                "key": key,  
                "latency": latency,  
                "key_states": key_states  
            }))  
                  
        elif action == 'press':  
            # 先释放所有按下的键  
            for other_key in key_states:  
                if key_states[other_key]:  
                    key_states[other_key] = False  
                    await asyncio.to_thread(keyboard.release, SUPPORTED_KEYS[other_key])  
                    print(f"[Direct Input] Auto-released key: {other_key}")  
              
            # 执行按下和释放操作  
            await asyncio.to_thread(keyboard.press, SUPPORTED_KEYS[key])  
            await asyncio.sleep(0.05)  # 短暂延迟，模拟真实按键  
            await asyncio.to_thread(keyboard.release, SUPPORTED_KEYS[key])  
            print(f"[Direct Input] Key Press: {key}")  
              
            latency = (time.time() - start_time) * 1000  
            await websocket.send(json.dumps({  
                "status": "success",  
                "action": "key_press",  
                "key": key,  
                "latency": latency,  
                "key_states": key_states  
            }))  
              
        elif action == 'get_state':  
            await websocket.send(json.dumps({  
                "status": "success",  
                "key_states": key_states  
            }))  
              
        elif action == 'reset_all':  
            for k in SUPPORTED_KEYS:  
                if key_states[k]:  
                    key_states[k] = False  
                    await asyncio.to_thread(keyboard.release, SUPPORTED_KEYS[k])  
                    print(f"[Direct Input] Reset key: {k}")  
              
            await websocket.send(json.dumps({  
                "status": "success",  
                "message": "All keys reset",  
                "key_states": key_states  
            }))  
              
        else:  
            await websocket.send(json.dumps({  
                "status": "error",  
                "message": f"Unknown action: {action}"  
            }))  
              
    except json.JSONDecodeError:  
        await websocket.send(json.dumps({  
            "status": "error",  
            "message": "Invalid JSON format"  
        }))  
    except Exception as e:  
        print(f"Error processing message: {str(e)}")  
        await websocket.send(json.dumps({  
            "status": "error",  
            "message": str(e)  
        }))  
  
# 广播按键状态给所有连接  
async def broadcast_key_states():  
    if active_connections:  
        message = json.dumps({  
            "type": "state_update",  
            "key_states": key_states  
        })  
        for conn in active_connections:  
            try:  
                await conn.send(message)  
            except Exception as e:  
                print(f"Error broadcasting to connection: {str(e)}")  
  
# 定期广播按键状态  
@app.before_serving  
async def setup_periodic_tasks():  
    app.broadcast_task = asyncio.create_task(periodic_broadcast())  
  
async def periodic_broadcast():  
    while True:  
        try:  
            await broadcast_key_states()  
        except Exception as e:  
            print(f"Error in periodic broadcast: {str(e)}")  
        await asyncio.sleep(1)  # 每秒广播一次  
  
@app.after_serving  
async def cleanup():  
    app.broadcast_task.cancel()  
    await app.broadcast_task  
  
# HTTP 路由用于健康检查  
@app.route("/health")  
async def health_check():  
    return {"status": "ok", "supported_keys": list(SUPPORTED_KEYS.keys())}  
  
if __name__ == "__main__":  
    print("Starting WebSocket server at ws://0.0.0.0:5000/ws")  
    print(f"Supported keys: {list(SUPPORTED_KEYS.keys())}")  
    app.run(host="0.0.0.0", port=5000)